package com.aliyun.dkms.gcs.sdk.example;

import com.aliyun.dkms.gcs.openapi.models.Config;
import com.aliyun.dkms.gcs.sdk.Client;
import com.aliyun.dkms.gcs.sdk.models.*;
import com.aliyun.tea.TeaException;

import java.io.Serializable;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;

/**
 * 这是使用KMS实例API <a href="https://help.aliyun.com/zh/kms/developer-reference/advancegeneratedatakey">AdvanceGenerateDataKey</a> 生成数据密钥和使用 <a href="https://help.aliyun.com/zh/kms/developer-reference/advancedecrypt">AdvanceDecrypt</a> 解密数据密钥密文的一个示例。
 * This is a sample to show how to use KMS Instance API <a href="https://help.aliyun.com/zh/kms/developer-reference/advancegeneratedatakey">AdvanceGenerateDataKey</a> to generate a data-key and use <a href="https://help.aliyun.com/zh/kms/developer-reference/advancedecrypt">AdvanceDecrypt</a> to decrypt the data-key from ciphertextBlob.
 */
public class AdvanceGenerateDataKeySample {

    // KMS实例客户端
    private static Client client = null;

    public static void main(String[] args) {
        try {
            // 1. 初始化KMS实例SDK. Init the KMS Instance SDK
            initClient();

            // 设置用于加密数据密钥的主密钥。注意：需要使用对称加密算法的密钥。
            // Set the keyId parameter. It will be used to encrypt the data-key generated by AdvanceGenerateDataKey, that is, the ciphertextBlob returned.
            // Note: The key should be a symmetric key.
            String keyId = "<your-key-id-or-alias>";
            Integer numberOfBytes = 32; // set data-key size (in bytes) on your need. For example, for AES-256, it is 32 bytes.

            // set the aad parameter on your need.
            byte[] aad = "<your-aad>".getBytes(StandardCharsets.UTF_8);
            // if you don't need to use aad , let it be null.
            // byte[] aad = null;
            // 2. Call AdvanceGenerateDataKey to generate data key
            AdvanceGenerateDataKeyResponse advanceGenerateDataKeyResponse = advanceGenerateDataKeySample(keyId, numberOfBytes, aad);
            byte[] dataKeyPlaintext = advanceGenerateDataKeyResponse.plaintext;

            // In this sample, do a check to make sure the dataKeyPlaintext is generated as expected.
            if (dataKeyPlaintext == null || dataKeyPlaintext.length != numberOfBytes) {
                System.err.println("It should not happen: dataKeyPlaintext is null or length not equals the setting for data-key size in bytes ");
                return;
            }
            // 3. You can use the plaintext data-key to encrypt your data locally.
            // <Your code for encrypting data with dataKeyPlaintext>

            // 4. Save data-key's ciphertextBlob and aad to storage.
            // When you need to get the plaintext data-key, you can read them from storage into an AdvanceGenerateDataKeyContext object and decrypt like below.
            AdvanceGenerateDataKeyContext dataKeyCiphertextContext = new AdvanceGenerateDataKeyContext(advanceGenerateDataKeyResponse.getCiphertextBlob(), aad);

            // 4. Call AdvanceDecrypt to decrypt data-key
            byte[] dataKeyPlaintext2 = decryptDataKeySample(dataKeyCiphertextContext);

            // In this sample, do a check to make sure the two plaintexts are equal.
            if (!Arrays.equals(dataKeyPlaintext2, dataKeyPlaintext)) {
                System.err.println("It should not happen: the two plaintexts are not equal.");
            }
            // 5. You can use data-key plaintext deciphered to decrypt data encrypted before locally.
            // < Your code for decrypting data with dataKeyPlaintext2 >
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 初始化KMS实例SDK。
     * Initialize the KMS Instance SDK.
     *
     * @throws Exception
     */
    public static void initClient() throws Exception {
        Config config = new Config();
        config.setProtocol("https"); // KMS instance uses the https protocol.

        config.setEndpoint("<your-endpoint>"); // 设置KMS实例Endpoint
        config.setClientKeyFile("<your-client-key-file>"); // 设置client key文件地址
        config.setPassword("<your-password>"); // 设置client key保护密码
        config.setCaFilePath("<path/to/yourCaCert>"); // 设置KMS实例的CA证书，可通过文件路径
        //config.setCa("<your-ca-certificate-content");// 或者，设置为服务端证书内容

        client = new Client(config);
    }

    /**
     * 调用AdvanceGenerateDataKey接口。
     *
     * @param keyId         设置接口的 KeyId 参数。（返回结果中的ciphertextBlob为使用该密钥加密的数据密钥密文。）
     * @param numberOfBytes 需生成的数据密钥的长度
     * @param aad           用户可根据使用的需要设置GCM加密模式的AAD参数。如果不需要使用，请设置为null或不设置。
     * @return AdvanceGenerateDataKeyResponse
     */
    private static AdvanceGenerateDataKeyResponse advanceGenerateDataKeySample(String keyId, Integer numberOfBytes, byte[] aad) {
        // 构建请求
        AdvanceGenerateDataKeyRequest request = new AdvanceGenerateDataKeyRequest();
        request.setKeyId(keyId);
        request.setNumberOfBytes(numberOfBytes);
        request.setAad(aad);
        try {
            AdvanceGenerateDataKeyResponse response = client.advanceGenerateDataKey(request);
            System.out.printf("RequestId: %s%n", response.getRequestId());
            return response;
        } catch (TeaException e) {
            System.out.printf("code: %s%n", e.getCode());
            System.out.printf("message: %s%n", e.getMessage());
            System.out.printf("requestId: %s%n", e.getData().get("requestId"));
            e.printStackTrace();
            throw new RuntimeException(e);
        } catch (Exception e) {
            System.out.printf("encrypt err: %s%n", e.getMessage());
            e.printStackTrace();
            throw new RuntimeException(e);
        }
    }

    /**
     * 调用AdvanceDecrypt接口解密通过AdvanceGenerateDataKey生成的数据密钥。
     *
     * @param advanceGenerateDataKeyContext 数据密钥密文和相关参数(aad)的封装。
     * @return 数据密钥的明文
     */
    private static byte[] decryptDataKeySample(AdvanceGenerateDataKeyContext advanceGenerateDataKeyContext) {
        // 构建AdvanceDecrypt请求
        AdvanceDecryptRequest request = new AdvanceDecryptRequest();
        request.setCiphertextBlob(advanceGenerateDataKeyContext.getCiphertextBlob());
        request.setAad(advanceGenerateDataKeyContext.getAad());

        try {
            AdvanceDecryptResponse response = client.advanceDecrypt(request);
            return response.plaintext;
        } catch (TeaException e) {
            System.out.printf("code: %s%n", e.getCode());
            System.out.printf("message: %s%n", e.getMessage());
            System.out.printf("requestId: %s%n", e.getData().get("requestId"));
            e.printStackTrace();
            throw new RuntimeException(e);
        } catch (Exception e) {
            System.out.printf("encrypt err: %s%n", e.getMessage());
            e.printStackTrace();
            throw new RuntimeException(e);
        }
    }

    /**
     * 保存AdvanceGenerateDataKey数据密钥密文和相关参数的示例。
     * A sample for storing AdvanceGenerateDataKey data key ciphertext and related parameters (aad).
     */
    static class AdvanceGenerateDataKeyContext implements Serializable {
        private byte[] ciphertextBlob;
        private byte[] aad;

        public AdvanceGenerateDataKeyContext() {
        }

        public AdvanceGenerateDataKeyContext(byte[] ciphertextBlob) {
            this.ciphertextBlob = ciphertextBlob;
            this.aad = null;
        }

        public AdvanceGenerateDataKeyContext(byte[] ciphertextBlob, byte[] aad) {
            this.ciphertextBlob = ciphertextBlob;
            this.aad = aad;
        }

        public byte[] getCiphertextBlob() {
            return ciphertextBlob;
        }

        public void setCiphertextBlob(byte[] ciphertextBlob) {
            this.ciphertextBlob = ciphertextBlob;
        }

        public byte[] getAad() {
            return aad;
        }

        public void setAad(byte[] aad) {
            this.aad = aad;
        }
    }
}
